// Copyright 2022 Huawei Cloud Computing Technology Co., Ltd.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include "CaeWebSocket.h"
#include "../cae_3rd/sha1/sha1.h"
using namespace std;

enum PARSE_STATE {
    NAME = 0,
    VALUE,
    SAW_CR,
    SKIPPING,
};

CaeWebSocket::CaeWebSocket(CaeTcpSocket *socket)
{
    m_tcpSocket = socket;
    m_webSocket = nullptr;
    m_sendBuf = nullptr;
    m_rxBuf = nullptr;
    m_isMaskZero = false;
    m_isMasked = false;
    memset_s(m_mask, sizeof(m_mask), 0, sizeof(m_mask));
    m_userBuf = nullptr;
    m_maskIdx = '\0';
    m_opcode = '\0';
    m_userBufLen = 0;
    m_wsState = 0;
    m_rxBufLen = 0;
    m_rxBufHead = nullptr;
    m_payloadLen = 0;
    m_parsingState = WS_NEW;

    pthread_mutex_init(&m_lock, nullptr);
    m_sendFlag = true;
}

CaeWebSocket::CaeWebSocket()
{
    m_tcpSocket = nullptr;
    m_webSocket = nullptr;
    m_sendBuf = nullptr;
    m_rxBuf = nullptr;
    m_isMaskZero = false;
    m_isMasked = false;
    memset_s(m_mask, sizeof(m_mask), 0, sizeof(m_mask));
    m_userBuf = nullptr;
    m_maskIdx = '\0';
    m_opcode = '\0';
    m_userBufLen = 0;
    m_wsState = 0;
    m_rxBufLen = 0;
    m_rxBufHead = nullptr;
    m_payloadLen = 0;
    m_parsingState = WS_NEW;
    pthread_mutex_init(&m_lock, nullptr);
    m_sendFlag = true;
}

bool CaeWebSocket::InitCaeWebSocket()
{
    m_sendBuf = (uint8_t *)malloc(MAX_PAYLOAD);
    if (m_sendBuf == nullptr) {
        LOGE("CaeWebSocket malloc m_sendBuf failed.");
        return false;
    }
    m_rxBuf = (uint8_t *)malloc(STREAM_MAX_BUF_LEN);
    if (m_rxBuf == nullptr) {
        LOGE("CaeWebSocket malloc m_rxBuf failed.");
        return false;
    }
    m_webSocket = new WebSocket();
    if (m_webSocket == nullptr) {
        LOGE("CaeWebSocket new m_webSocket failed.");
        return false;
    }
    return true;
}

CaeWebSocket::~CaeWebSocket()
{
    if (m_sendBuf != nullptr) {
        free(m_sendBuf);
        m_sendBuf = nullptr;
    }

    if (m_webSocket != nullptr) {
        delete m_webSocket;
        m_webSocket = nullptr;
    }

    if (m_rxBuf != nullptr) {
        free(m_rxBuf);
        m_rxBuf = nullptr;
    }

    if (m_tcpSocket != nullptr) {
        delete m_tcpSocket;
        m_tcpSocket = nullptr;
    }

    pthread_mutex_destroy(&m_lock);
}

bool CaeWebSocket::ParseHttpHeader()
{
    uint32_t limit = 4096;
    uint32_t totalLen = 0;
    WebSocketFrameType frameType;
    while (totalLen <= limit && m_rxBuf != nullptr) {
        int ret;
        if (m_tcpSocket != nullptr && m_tcpSocket->GetStatus() == SOCKET_STATUS_NORMAL) {
            ret = m_tcpSocket->Recv(m_rxBuf + totalLen, limit - totalLen);
        } else {
            LOGE("Tcp socket's status is error.");
            return false;
        }
        if (ret <= 0) {
            LOGE("Socket disconnected when parsing http header.");
            return false;
        }
        totalLen += static_cast<uint32_t>(ret);
        if (m_webSocket != nullptr) {
            frameType = m_webSocket->parseHandshake(m_rxBuf, static_cast<int>(totalLen));
        } else {
            LOGE("Web socket is null.");
            return false;
        }
        if (frameType == INCOMPLETE_FRAME) {
            continue;
        } else {
            return true;
        }
    }
    return false;
}

bool CaeWebSocket::Handshake()
{
    LOGI("Start parse http header.");
    if (ParseHttpHeader()) {
        string answer = m_webSocket->answerHandshake();
        int sendSize;
        if (m_tcpSocket != nullptr && m_tcpSocket->GetStatus() == SOCKET_STATUS_NORMAL) {
            sendSize = m_tcpSocket->Send((void *)answer.c_str(), answer.size());
        } else {
            LOGE("Tcp socket's status is error.");
            return false;
        }
        if (sendSize != static_cast<int>(answer.size())) {
            LOGE("Failed to Send response.");
            return false;
        } else {
            return true;
        }
    } else {
        LOGE("Failed to parsing http header.");
        return false;
    }
}

int CaeWebSocket::Send(void *uBuf, size_t len)
{
    int ret = 0;
    pthread_mutex_lock(&m_lock);
    if (m_sendFlag) {
        ret = Send(uBuf, len, WB_OPC_BINARY_FRAME);
    }
    pthread_mutex_unlock(&m_lock);
    return ret;
}

int CaeWebSocket::SendCloseResponse()
{
    return Send(nullptr, 0, WB_OPC_CLOSE);
}

int CaeWebSocket::SendClose()
{
    return Send(nullptr, 0, WB_OPC_CLOSE);
}

void CaeWebSocket::EnableSend()
{
    pthread_mutex_lock(&m_lock);
    m_sendFlag = true;
    pthread_mutex_unlock(&m_lock);
}

void CaeWebSocket::DisableSend()
{
    pthread_mutex_lock(&m_lock);
    m_sendFlag = false;
    pthread_mutex_unlock(&m_lock);
}

int CaeWebSocket::Send(void *uBuf, size_t len, enum WB_OPCODE wp)
{
    int ret = 0;
    if (len > MAX_PAYLOAD) {
        LOGE("Failed to Send buffer, because len > MAX_PAYLOAD.");
        return 0;
    }

    if (m_tcpSocket == nullptr) {
        LOGE("Failed to Send buffer, tcp socket is null.");
        return 0;
    }

    uint32_t frameType = wp;
    frameType |= 1 << 7;
    int size = m_webSocket->makeFrame((WebSocketFrameType)frameType, (uint8_t *)uBuf, (int)len, m_sendBuf,
                                      MAX_PAYLOAD);
    if (m_tcpSocket->GetStatus() == SOCKET_STATUS_NORMAL) {
        ret = m_tcpSocket->Send((uint8_t *)m_sendBuf, size);
    }
    return ret;
}

int CaeWebSocket::Spill()
{
    if (m_opcode == WB_OPC_CLOSE) {
        LOGD("Server sees client close packet.");
        m_wsState = 0;
        DisableSend();
        SendCloseResponse();
        this->EventNotice();
        return SOCKET_RECV_FAIL_DISCONN;
    }
    return 0;
}

int CaeWebSocket::HandleWsFrameHdrLen(uint8_t c)
{
    int ret = 0;
    m_isMasked = ((c & 0x80) > 0);

    switch (c & 0x7f) {
        case 126:
            if (m_opcode & 8) {
                LOGE("Control frame length is illegal.");
                return -1;
            }
            m_parsingState = WS_FRAME_HDR_LEN16_2;
            break;
        case 127:
            if (m_opcode & 8) {
                LOGE("Control frame length is illegal.");
                return -1;
            }
            m_parsingState = WS_FRAME_HDR_LEN64_8;
            break;
        default:
            m_payloadLen = c & 0x7f;
            if (m_isMasked) {
                m_parsingState = WS_COLLECT_FRAME_KEY_1;
            } else {
                if (m_payloadLen) {
                    m_parsingState = WS_PAYLOAD;
                } else {
                    m_parsingState = WS_NEW;
                    ret = Spill();
                }
            }
            break;
    }
    return ret;
}

int CaeWebSocket::HandleWsFrameHdr1(uint8_t c)
{
    m_opcode = c & 0xf;
    switch (m_opcode) {
        case WB_OPC_TEXT_FRAME:
        case WB_OPC_BINARY_FRAME:
            break;
        case 3:
        case 4:
        case 5:
        case 6:
        case 7:
        case 0xb:
        case 0xc:
        case 0xd:
        case 0xe:
        case 0xf:
            LOGE("websocket illegal opcode = %x", c);
            return -1;
    }
    m_parsingState = WS_FRAME_HDR_LEN;
    return 0;
}
int CaeWebSocket::ParseHeader(uint8_t c)
{
    int ret = 0;

    switch (m_parsingState) {
        case WS_NEW:
            m_isMaskZero = true;

        case WS_FRAME_HDR_1:
            ret = HandleWsFrameHdr1(c);
            break;

        case WS_FRAME_HDR_LEN:
            ret = HandleWsFrameHdrLen(c);
            break;

        case WS_FRAME_HDR_LEN16_2:
            m_payloadLen = static_cast<size_t>(c) << 8;
            m_parsingState = WS_FRAME_HDR_LEN16_1;
            break;

        case WS_FRAME_HDR_LEN16_1:
            m_payloadLen |= c;
            if (m_isMasked) {
                m_parsingState = WS_COLLECT_FRAME_KEY_1;
            } else {
                m_parsingState = WS_PAYLOAD;
            }
            break;

        case WS_FRAME_HDR_LEN64_8:
            if (c & 0x80) {
                LOGE("Error: length need to = 0");
                return -1;
            }
#if defined __LP64__
            m_payloadLen = ((size_t)c) << 56;
#else
            m_payload_len = 0;
#endif
            m_parsingState = WS_FRAME_HDR_LEN64_7;
            break;

        case WS_FRAME_HDR_LEN64_7:
#if defined __LP64__
            m_payloadLen |= ((size_t)c) << 48;
#endif
            m_parsingState = WS_FRAME_HDR_LEN64_6;
            break;

        case WS_FRAME_HDR_LEN64_6:
#if defined __LP64__
            m_payloadLen |= ((size_t)c) << 40;
#endif
            m_parsingState = WS_FRAME_HDR_LEN64_5;
            break;

        case WS_FRAME_HDR_LEN64_5:
#if defined __LP64__
            m_payloadLen |= ((size_t)c) << 32;
#endif
            m_parsingState = WS_FRAME_HDR_LEN64_4;
            break;

        case WS_FRAME_HDR_LEN64_4:
            m_payloadLen |= ((size_t)c) << 24;
            m_parsingState = WS_FRAME_HDR_LEN64_3;
            break;

        case WS_FRAME_HDR_LEN64_3:
            m_payloadLen |= ((size_t)c) << 16;
            m_parsingState = WS_FRAME_HDR_LEN64_2;
            break;

        case WS_FRAME_HDR_LEN64_2:
            m_payloadLen |= ((size_t)c) << 8;
            m_parsingState = WS_FRAME_HDR_LEN64_1;
            break;

        case WS_FRAME_HDR_LEN64_1:
            m_payloadLen |= ((size_t)c);
            if (m_isMasked) {
                m_parsingState = WS_COLLECT_FRAME_KEY_1;
            } else {
                m_parsingState = WS_PAYLOAD;
            }
            break;
        default:
            ret = HandleWsCollectFrameKey(c);
            break;
    }

    return ret;
}

int CaeWebSocket::HandleWsCollectFrameKey(uint8_t c)
{
    int ret = 0;
    switch (m_parsingState) {
        case WS_COLLECT_FRAME_KEY_1:
            m_mask[0] = c;
            m_parsingState = WS_COLLECT_FRAME_KEY_2;
            break;

        case WS_COLLECT_FRAME_KEY_2:
            m_mask[1] = c;
            m_parsingState = WS_COLLECT_FRAME_KEY_3;
            break;

        case WS_COLLECT_FRAME_KEY_3:
            m_mask[2] = c;
            m_parsingState = WS_COLLECT_FRAME_KEY_4;
            break;

        case WS_COLLECT_FRAME_KEY_4:
            m_mask[3] = c;
            if (m_mask[0] > 0 || m_mask[1] > 0 || m_mask[2] > 0 || m_mask[3] > 0) {
                m_isMaskZero = false;
            }

            m_parsingState = WS_PAYLOAD;
            m_maskIdx = 0;
            if (m_payloadLen == 0) {
                m_parsingState = WS_NEW;
                ret = Spill();
            }
            break;
        default:
            break;
    }
    return ret;
}

int CaeWebSocket::ReadPayload()
{
    uint8_t mask[4];

    uint32_t minSize = m_userBufLen;

    if (minSize > m_payloadLen) {
        minSize = m_payloadLen;
    }

    if (minSize > m_rxBufLen) {
        minSize = m_rxBufLen;
    }

    if ((m_opcode == WB_OPC_BINARY_FRAME || m_opcode == WB_OPC_CONTINUATION || m_opcode == WB_OPC_TEXT_FRAME ||
         m_opcode == WB_OPC_CLOSE) &&
        m_rxBufHead != nullptr) {
        if (m_isMaskZero) {
            int ret = memcpy_s(m_userBuf, m_userBufLen, m_rxBufHead, minSize);
            if (ret != EOK) {
                LOGE("Failed to memcpy_s m_userBuf.");
            }
            m_rxBufHead += minSize;
            m_userBuf += minSize;
            m_userBufLen -= minSize;
            m_payloadLen -= minSize;
            m_rxBufLen -= minSize;
        } else {
            uint32_t i;
            for (i = 0; i < 4; ++i) {
                mask[i] = m_mask[(m_maskIdx + i) & 3];
            }

            i = minSize / 4;
            while (i--) {
                *(m_userBuf++) = *(m_rxBufHead++) ^ mask[0];
                *(m_userBuf++) = *(m_rxBufHead++) ^ mask[1];
                *(m_userBuf++) = *(m_rxBufHead++) ^ mask[2];
                *(m_userBuf++) = *(m_rxBufHead++) ^ mask[3];
            }

            for (i = 0; i < (minSize & 3); ++i) {
                *(m_userBuf++) = *(m_rxBufHead++) ^ mask[i];
            }

            m_maskIdx = (m_maskIdx + minSize) & 3;

            m_userBufLen -= minSize;
            m_payloadLen -= minSize;
            m_rxBufLen -= minSize;
        }
    }
    return static_cast<int>(minSize);
}

int CaeWebSocket::Recv(void *buf, size_t len)
{
    if (buf == nullptr || len <= 0) {
        LOGE("Invalid parameter in Recv.");
        return 0;
    }

    m_userBuf = (uint8_t *)buf;
    m_userBufLen = len;

    int ret = -1;
    if (m_rxBufLen == 0) {
        if (m_tcpSocket != nullptr && m_tcpSocket->GetStatus() == SOCKET_STATUS_NORMAL) {
            ret = m_tcpSocket->Recv(m_rxBuf, STREAM_MAX_BUF_LEN);
        }
        if (ret <= 0) {
            return ret;
        } else {
            m_rxBufLen = static_cast<size_t>(ret);
        }
        m_rxBufHead = m_rxBuf;
    }

    int minSize = static_cast<int>((m_rxBufLen < len) ? m_rxBufLen : len);
    int m;
    int readLen = 0;
    while (minSize) {
        if (m_parsingState == WS_PAYLOAD) {
            m = ReadPayload();
            readLen += m;
            minSize -= m;
            if (m_payloadLen == 0) {
                m_parsingState = WS_NEW;
                if (m_opcode == WB_OPC_CLOSE) {
                    LOGE("Client request to disconnect.");
                    DisableSend();
                    SendCloseResponse();
                    this->EventNotice();
                    LOGE("Client request to disconnect out.");
                    return SOCKET_RECV_FAIL_DISCONN;
                }
            }
        }

        if (minSize <= 0) {
            break;
        }

        m = ParseHeader(*(m_rxBufHead)++);
        minSize--;
        m_rxBufLen--;
        if (m < 0) {
            LOGE("Web socket error in ParseHeader.");
            return SOCKET_RECV_FAIL_DISCONN;
        }
    }
    return readLen;
}

int CaeWebSocket::GetFd()
{
    if (m_tcpSocket) {
        return m_tcpSocket->GetFd();
    } else {
        return -1;
    }
}

void CaeWebSocket::SetSocketOption(uint32_t localSocketOption)
{
    if (m_tcpSocket) {
        m_tcpSocket->SetSocketOption(localSocketOption);
    }
}

void CaeWebSocket::CloseSocket()
{
    if (m_tcpSocket != nullptr) {
        m_tcpSocket->CloseSocket();
    }
}
